#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <gmssl/sm2.h>
#include <gmssl/sm3.h>
#include <gmssl/sm4.h>
#include <gmssl/rand.h>

void print_hex(const char *label, const uint8_t *data, size_t len) {
    printf("%s: ", label);
    for (size_t i = 0; i < len; i++) {
        printf("%02X", data[i]);
    }
    printf("\n");
}

int main() {
    int ret = 0;

    // Step 1: 生成接收方 SM2 密钥对并保存
    SM2_KEY receiver_key;
    sm2_key_generate(&receiver_key);
    FILE *key_file = fopen("receiver_key.pem", "wb");
    sm2_private_key_to_pem(&receiver_key, key_file); // GmSSL 3.1.1 兼容函数
    fclose(key_file);
    printf("[发送方] 接收方密钥对已保存至 receiver_key.pem\n");

    // Step 2: 生成随机 SM4 对称密钥
    uint8_t sym_key[16];
    rand_bytes(sym_key, sizeof(sym_key));
    print_hex("[发送方] SM4 对称密钥", sym_key, sizeof(sym_key));

    // Step 3: 准备明文并填充（PKCS#7）
    const char *plaintext = "Hello, 数字信封!";
    size_t plaintext_len = strlen(plaintext);
    size_t block_size = 16; // SM4 分组大小
    size_t pad_value = block_size - (plaintext_len % block_size);
    size_t padded_len = plaintext_len + pad_value;
    uint8_t *padded_plaintext = malloc(padded_len);
    memcpy(padded_plaintext, plaintext, plaintext_len);
    memset(padded_plaintext + plaintext_len, pad_value, pad_value);

    // Step 4: 使用 SM4-CBC 加密
    uint8_t iv[16];
    rand_bytes(iv, sizeof(iv));
    SM4_KEY sm4_enc_key;
    sm4_set_encrypt_key(&sm4_enc_key, sym_key);
    uint8_t *ciphertext = malloc(padded_len);
    sm4_cbc_encrypt(&sm4_enc_key, iv, padded_plaintext, padded_len, ciphertext);
    print_hex("[发送方] 密文", ciphertext, padded_len);

    // Step 5: 计算明文哈希
    uint8_t hash[SM3_DIGEST_SIZE];
    SM3_CTX sm3_ctx;
    sm3_init(&sm3_ctx);
    sm3_update(&sm3_ctx, (uint8_t *)plaintext, plaintext_len);
    sm3_finish(&sm3_ctx, hash);
    print_hex("[发送方] 明文哈希", hash, SM3_DIGEST_SIZE);

    // Step 6: 加密对称密钥
    uint8_t encrypted_key[SM2_MAX_CIPHERTEXT_SIZE];
    size_t encrypted_key_len;
    ret = sm2_encrypt(&receiver_key, sym_key, sizeof(sym_key), encrypted_key, &encrypted_key_len);
    if (ret != 1) {
        fprintf(stderr, "[发送方] SM2 加密失败\n");
        return EXIT_FAILURE;
    }
    print_hex("[发送方] 加密后的对称密钥", encrypted_key, encrypted_key_len);

    // Step 7: 保存数字信封
    FILE *envelope_file = fopen("envelope.bin", "wb");
    fwrite(&encrypted_key_len, sizeof(size_t), 1, envelope_file);
    fwrite(encrypted_key, 1, encrypted_key_len, envelope_file);
    fwrite(iv, 1, sizeof(iv), envelope_file);
    fwrite(&padded_len, sizeof(size_t), 1, envelope_file);
    fwrite(ciphertext, 1, padded_len, envelope_file);
    fwrite(hash, 1, SM3_DIGEST_SIZE, envelope_file);
    fclose(envelope_file);
    printf("[发送方] 数字信封已保存至 envelope.bin\n");

    // 释放内存
    free(padded_plaintext);
    free(ciphertext);
    return EXIT_SUCCESS;
}
